#!/usr/bin/env python

from optparse import OptionParser
import os
import os.path
import sys

import charmicat
import histogram


def main():
    op = OptionParser(usage='%prog [opts] file.pkl')


    op.add_option("-f", "--force", action="store_true", dest="force", 
                  default=False, help="Force overwrite of output directory")
    op.add_option("-n", "--name", action='append', dest='names', default=[],
                  help='Write images only for histogram NAME (can be repeated)',
                  metavar='NAME')
    op.add_option("-o", "--output-dir", dest="output", default=None, 
                  help="Place output in directory DIR", metavar="DIR")
    op.add_option("-p", "--presentation", action="store_true", 
                  dest="presentation", default=False, 
                  help="Draw images in presentation mode")
    op.add_option("-t", "--tag", dest='tag', default='',
                  help='Append TAG to all histogram titles',
                  metavar='TAG')
    op.add_option("--cmap", dest='cmap', default='hot', 
                  help='Specify colormap CMAP for 2D hists', metavar='CMAP')
    op.add_option("-x", "--ext", dest='ext', default=[], action='append',
                  help='Write out files of type EXT (can be repeated)', 
                  metavar='EXT')
    op.add_option("--title", dest='title', default=None,
                  help='Set all histogram titles to TITLE', metavar='TITLE')

    opts, args = op.parse_args()

    if len(args) != 1:
        sys.stderr.write("Please supply *one* pickle file with histograms.\n")
        sys.stderr.write("See `draw_hists.py -h' for more.\n")
        sys.stderr.flush()
        return -1

    infname = args[0]

    if opts.output is None:
        opts.output = infname.rsplit('.', 1)[0]

    pngdir = os.path.join(opts.output, 'png')
    pdfdir = os.path.join(opts.output, 'pdf')

    opts.output = charmicat.safe_mkdir(
        opts.output, default=infname.rsplit('.',1)[0], force=opts.force)

    opts.fontsize = 'medium'
    if opts.presentation:
        opts.fontsize = 'xx-large'

    if len(opts.ext) == 0:
        opts.ext = 'png'
    else:
        for i in xrange(len(opts.ext)):
            if opts.ext[i] == '.':
                opts.ext[i] = opts.ext[1:]
    
    hc = histogram.load(infname, names=opts.names)

    if type(hc) in histogram.htypes:
        h = hc
        hc = histogram.HistogramCollection()
        hc.add(h)

    nhists = len(hc)
    ihist = 0
    for h in hc.iter1d():
        sys.stdout.write("{}/{} hists complete\r".format(ihist, nhists))
        sys.stdout.flush()

        out = os.path.join(opts.output, h.name)

        if len(h) == 0:
            with open("{0}_is_empty".format(out), 'w') as fp:
                pass
            
            continue

        if opts.title is None:
            opts.title = h._plotattrs['title']

        h.image_prepare(fontsize=opts.fontsize, tag=opts.tag, 
                        title=opts.title)
        h.image_save(out=out, ext=opts.ext)
        ihist += 1

    for h in hc.iter2d():        
        sys.stdout.write("{}/{} hists complete\r".format(ihist, nhists))
        sys.stdout.flush()

        out = os.path.join(opts.output, h.name)

        if len(h) == 0:
            with open("{0}_is_empty".format(out), 'w') as fp:
                pass            
            continue
        
        h.image_prepare(fontsize=opts.fontsize,
                        tag=opts.tag, cmap=opts.cmap, title=opts.title)
        h.image_save(out=out, ext=opts.ext)
        
        ihist += 1

    print "{0}/{0} hists complete".format(nhists)






if __name__ == "__main__":
    try:
        sys.exit(main())
    except charmicat.BadCommandLineError, e:
        import os
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `{0} -h' for more.\n".format(os.path.basename(sys.argv[0])))
        sys.stderr.flush()
        sys.exit(e.retcode)
